# -*- coding: utf-8 -*-
import matplotlib.pyplot as plt
import numpy as np


def plot_logfile(file_path, metric1='Train Loss',metric2='Valid Loss', show_fig=True, save_pic=False, save_pic_path=None):
    # 读取文件并解析数据
    log_data = []

    # 由于输入的metric1和metric2在log文件中可能不存在，因此我们需要做判断
    metric1_exists = False
    metric2_exists = False

    with open(file_path, 'r') as file:
        for line in file:
            if 'Epoch' in line and (metric1 in line or metric2 in line):
                parts = line.split('; ')
                re_dict ={}
                for part in parts:
                    if 'Epoch' in part:
                        epoch = int(part.split(': ')[-1].strip())
                        re_dict['Epoch'] = epoch
                    elif metric1 in part:
                        metric1_exists = True
                        metric1_vals = float(part.split(': ')[1].strip())
                        re_dict[metric1] = metric1_vals
                    elif metric2 in part:
                        metric2_exists = True
                        metric2_vals = float(part.split(': ')[1].strip())
                        re_dict[metric2] = metric2_vals
                log_data.append(re_dict)
    if metric1_exists==False and metric2_exists==False:
        raise ValueError(f"{metric1}和{metric2}在{file_path}文件中不存在！")

    # 提取数据到数组
    epochs = np.array([entry['Epoch'] for entry in log_data])

    # 绘制折线图
    plt.figure(figsize=(14, 7))
    # 第1个指标的折线图
    if metric1_exists:
        m1_vals = np.array([entry[metric1] for entry in log_data])
        plt.subplot(1, 2, 1)
        plt.plot(epochs, m1_vals, marker='o', linestyle='-', color='blue')
        plt.title(f'{metric1} over Epochs')
        plt.xlabel('Epoch')
        plt.ylabel(metric1)
        plt.grid(True)

    # 第2个指标的折线图
    if metric2_exists:
        m2_vals = np.array([entry[metric2] for entry in log_data])
        plt.subplot(1, 2, 2)
        plt.plot(epochs, m2_vals, marker='s', linestyle='-', color='green')
        plt.title(f'{metric2} over Epochs')
        plt.xlabel('Epoch')
        plt.ylabel(metric2)
        plt.grid(True)

    # 显示图表
    plt.tight_layout()
    if save_pic:
        plt.savefig(save_pic_path, dpi=300, bbox_inches='tight')
    if show_fig:
        plt.show()



if __name__ == '__main__':
    file_path = r'C:\tmp\gnnwr20240628_175932\20240628_175932.log'
    # plot_logfile(file_path, metric1='Train AIC',metric2='Valid R2')
    # plot_logfile(file_path, metric1='Train F1',metric2='Train Recall')
    plot_logfile(file_path, metric1='Train Loss', metric2='Valid Loss')
    # plot_logfile(file_path, metric1='Learning Rate', metric2='Valid Loss')